import json
import os
import base64
import cv2
import numpy as np

json_save_path = '/media/wsl/SB@data/instance_segmentation/大米/all/json'
txt_path = '/media/wsl/SB@data/instance_segmentation/大米/all/txt'  # psenet输出的txt
img_path = '/media/wsl/SB@data/instance_segmentation/大米/all/img'
null = 0


def txt2json(txt_path, json_save_path, img_path):
    """
    :param txt_path: psenet输出的txt
    :param json_save_path: 生成的json文件的保存路径
    :param img_path:图片的路径
    :return:
    """
    
    for txt_file in os.listdir(txt_path):
        base_name = txt_file.split('.')[0]
        
        if txt_file.split('.')[-1] == 'txt':
            # 读取txt里面的坐标信息
            txt_f = open(os.path.join(txt_path, txt_file), 'r', encoding='utf-8')
            txt_infor = [i.rstrip() for i in txt_f.readlines()]
            
            # 读取图片base64信息
            img_name = base_name + '.png'
            # print(os.path.join(img_path, img_name))
            h, w, c = cv2.imread(os.path.join(img_path, img_name)).shape
            # 获取图片宽高
            
            img_read = open(os.path.join(img_path, img_name), "rb").read()
            image_base64 = base64.b64encode(img_read)
            
            # 构建labelme字典
            json_dict = {
                "version": "4.2.9",
                "flags": {},
                "shapes": [],
                "imagePath": img_name,
                "imageData": str(image_base64, 'utf8'),
                "imageHeight": h,
                "imageWidth": w,
            }
            for infor in txt_infor:
                plist = infor.replace('[','').replace(']','').replace('  ','').split(',')
                plist=list(map(float,plist))
                
                # print(infor)
                plist=np.array(plist).reshape(-1,2)
                plist=plist.tolist()
                #plist=plist.astype(np.float32)
                
                label = '0'
                shape_dict = {
                    "label": label,
                    "points": plist,
                    "group_id": null,
                    "shape_type": "polygon",
                    "flags": {}
                }
                json_dict["shapes"].append(shape_dict)
            
            # 生成json文件
            json_path = os.path.join(json_save_path, base_name + '.json')
            
            json_file = open(json_path, 'w')
            json.dump(json_dict, json_file, indent=2)


if __name__ == '__main__':
    txt2json(txt_path, json_save_path, img_path)
